# task_executor.py
import traceback,time
from task_flow.serialization import xloads, xdumps
from task_flow.redis_locks import RedisLockWithWatchdog
from logging import getLogger

logger = getLogger(__name__)

__all__ = ['execute_task']

def execute_task(
        req_json: dict,
        node_info: dict,
        r,  # 传入已初始化的 Redis 客户端 (r)
        task_map,  # 传入该 worker 的 task_map
):
    """
    通用的 Celery 任务执行器，处理缓存、锁、执行和结果存储。

    Args:
        req_json (dict): 任务请求数据。
        node_info (dict): 包含单个任务配置的字典 {node_name: config}。
        r: 已初始化的 Redis 客户端实例。
        task_map (dict): 任务名称到任务对象实例的映射。

    Returns:
        dict: 包含 'code', 'data', 'msg' 的结果字典。
    """
    order_id = None
    node_name, node_conf = next(iter(node_info.items()))
    try:

        generator = task_map[node_name]
        base_info = req_json["base_info"]
        order_id = base_info["order_id"]
        expire = node_conf.get("expire", 15 * 60)
        cache_field = node_conf.get("cache_key", 'order_id')

        cache_risk_key = base_info[cache_field]
        logger.info(
            f"node_name:{node_name},order_id:{order_id},cache:{cache_field}:{cache_risk_key},接收到计算请求")
        save_key = f"{node_name}:{cache_field}:{cache_risk_key}"
        cached_data = r.get(save_key)
        if cached_data:
            logger.info(
                f"node_name:{node_name},order_id:{order_id},cache:{cache_field}:{cache_risk_key},返回缓存数据")
            return {
                "code": 200,
                "data": xloads(cached_data),
                "msg": ""
            }

        lock_key = f"lock:{save_key}"
        with RedisLockWithWatchdog(r, lock_key):
            cached_data = r.get(save_key)
            if cached_data:
                logger.info(
                    f"node_name:{node_name},order_id:{order_id},cache:{cache_field}:{cache_risk_key},返回缓存数据")
                return {
                    "code": 200,
                    "data": xloads(cached_data),
                    "msg": ""
                }

            t1 = time.time()
            data = generator.gen_features(req_json)
            t2 = time.time()

            logger.info(
                f"node_name:{node_name},order_id:{order_id},cache:{cache_field}:{cache_risk_key},完整计算耗时:{round(t2 - t1, 4)}")
            r.set(save_key, xdumps(data), ex=expire)
            logger.info(f"node_name:{node_name},order_id:{order_id},cache:{cache_field}:{cache_risk_key},写入缓存")
            return {
                "code": 200,
                "data": data,
                "msg": ""
            }
    except:
        erro_msg = f"node_name:{node_name},order_id:{order_id}计算失败,{traceback.format_exc()}"
        logger.exception(erro_msg)
        return {
            "code": -200,
            "msg": erro_msg,
            "data": None
        }
